from __future__ import print_function, division
import shutil
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
from torchvision import models
import cv2
import numpy as np
import os
from lib.network import Network
import os
import math
import torchvision


def vis_nl_map(img_path, nl_map, bt_idx, vis_size=(560, 560)):
    img = cv2.imread(img_path, 1)
    img = cv2.resize(img, dsize=vis_size)
    h, w, c = img.shape
    nl_map_0 = nl_map[0]
    total_region, nl_map_length = nl_map_0.shape
    region_per_row = round(math.sqrt(total_region))
    size_of_region = round(w / region_per_row)

    nl_map_size = round(math.sqrt(nl_map_length))

    for index in range(total_region):
        img_draw = img.copy()
        nl_map = nl_map_0[index]
        nl_map = nl_map.reshape(nl_map_size, nl_map_size)
        nl_map = cv2.resize(nl_map, dsize=(h, w))
        nl_map = np.uint8(nl_map * 255)
        heat_img = cv2.applyColorMap(nl_map, cv2.COLORMAP_JET)
        heat_img = cv2.cvtColor(heat_img, cv2.COLOR_BGR2RGB)
        img_add = cv2.addWeighted(img_draw, 0.3, heat_img, 0.7, 0)
        x0 = index // region_per_row * size_of_region
        x1 = x0 + size_of_region
        y0 = index % region_per_row * size_of_region
        y1 = y0 + size_of_region
        cv2.rectangle(img_add, (y0, x0), (y1, x1), (255, 0, 0), 1)
        cv2.imshow("nlp", img_add)
        cv2.imwrite("./nlp/%d_%d.jpg"%(bt_idx, index), img_add)
        cv2.waitKey(100)


def del_dir(path):
    for i in os.listdir(path):
        path_file = os.path.join(path, i)       # 取文件绝对路径
        if os.path.isfile(path_file):
            os.remove(path_file)
        else:
            del_dir(path_file)


def rm_mkdir(dir_path):
    if os.path.exists(dir_path):
        del_dir(dir_path)
        print('Clean path - %s' % dir_path)
    else:
        os.makedirs(dir_path)
        print('Create path - %s' % dir_path)


def load_model(test_dir):
    # test_datasets = torchvision.datasets.MNIST(root='./mnist/',
    #                                        transform=torchvision.transforms.ToTensor(),
    #                                        train=False)
    val_transforms = transforms.Compose([
        transforms.Grayscale(),
        transforms.ToTensor()
    ])
    test_datasets = datasets.ImageFolder(test_dir, transform=val_transforms)

    net = Network()
    if torch.cuda.is_available():
        net = nn.DataParallel(net)
        net.cuda()
    state_dict = torch.load(r".\mnist_net.pth")
    net.load_state_dict(state_dict)
    return net, test_datasets


def run():
    test_dir = r'H:\yuanbaoxi\ybx_gitee\non_local\mnist\testimgs'
    net, test_datasets = load_model(test_dir)
    test_dataloader = torch.utils.data.DataLoader(test_datasets, batch_size=1, shuffle=False)
    rm_mkdir("./result")
    name_list_en = os.listdir(test_dir)
    with torch.no_grad():
        correct = 0
        total = 0
        err_idx = 0
        cls_idx = test_datasets.class_to_idx
        cls = test_datasets.classes
        class_num = len(cls)
        conf_mat = np.zeros((class_num, class_num), np.uint32)
        bt_idx = 0
        for data in test_dataloader:
            net.eval()
            inputs_ori, labels_ori = data
            inputs, labels = Variable(inputs_ori).cuda(), Variable(labels_ori).cuda()
            _, nl_mep_list = net.module.forward_with_nl_map(inputs)
            # (b, h1*w1, h2*w2)
            nl_map_1 = nl_mep_list[0].cpu().numpy()
            nl_map_2 = nl_mep_list[1].cpu().numpy()

            img = torchvision.transforms.ToPILImage()(inputs.cpu()[0])
            img.save('nl_map_vis/sample.png')
            # np.save('nl_map_vis/nl_map_1', nl_map_1)
            # np.save('nl_map_vis/nl_map_2', nl_map_2)
            vis_nl_map(img_path='nl_map_vis/sample.png', nl_map=nl_map_1, bt_idx=bt_idx, vis_size=(560, 560))
            # vis_nl_map(img_path='nl_map_vis/sample.png', nl_map=nl_map_2, bt_idx=bt_idx, vis_size=(560, 560))

            bt_idx += 1
            # np.save('nl_map_vis/nl_map_1_%d'%bt_idx, nl_map_1)
            # np.save('nl_map_vis/nl_map_2_%d'%bt_idx, nl_map_2)


if __name__ == "__main__":
    run()
    # train()